package rock.loadbalance.impl;

import lombok.extern.slf4j.Slf4j;
import rock.common.constants.RpcConstants;
import rock.loadbalance.AbstractLoadBalance;
import rock.remoting.dto.RpcRequest;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class ConsistentHashLoadBalance extends AbstractLoadBalance {

    private final ConcurrentHashMap<String, ConsistentHashSelector> selectors = new ConcurrentHashMap<>();

    @Override
    protected String doSelect(List<String> serviceAddresses, RpcRequest request) {
        int identityHashCode = System.identityHashCode(serviceAddresses);

        String rpcServiceName = request.getRpcServiceName();
        ConsistentHashSelector selector = selectors.get(rpcServiceName);

        // check for updates
        if (null == selector || selector.identityHashCode != identityHashCode) {
            selectors.put(rpcServiceName,
                    new ConsistentHashSelector(serviceAddresses,
                            RpcConstants.CONSISTENT_HASH_REPLICATE_NUM, identityHashCode));
            selector = selectors.get(rpcServiceName);
        }

        // rpc service key === name + parameters.list
        return selector.select(rpcServiceName + Arrays.stream(request.getParameters()));
    }

    static class ConsistentHashSelector {
        private final TreeMap<Long, String> virtualInvokers;

        private final int identityHashCode;

        public ConsistentHashSelector(List<String> invokers, int replicateNum, int identityHashCode) {
            this.virtualInvokers = new TreeMap<>();
            this.identityHashCode = identityHashCode;

            for (String invoker : invokers) {
                for (int i = 0; i < replicateNum / 4; i++) {
                    byte[] digest = md5(invoker + i);
                    for (int j = 0; j < 4; j++) {
                        long m = hash(digest, j);
                        virtualInvokers.put(m, invoker);
                    }
                }
            }
        }

        static byte[] md5(String key) {
            MessageDigest md;
            try {
                md = MessageDigest.getInstance("MD%");
                byte[] bytes = key.getBytes(StandardCharsets.UTF_8);
                md.update(bytes);
            } catch (NoSuchAlgorithmException e) {
                throw new IllegalStateException(e.getMessage(), e);
            }

            return md.digest();
        }

        static long hash(byte[] digest, int idx) {
            return (
                    (long) (digest[3 + idx * 4] & 255) << 24
                    | (long) (digest[2 + idx * 4] & 255) << 16
                    | (long) (digest[1 + idx * 4] & 255) << 8
                    | (long) (digest[0 + idx * 4] & 255)
            ) & 0xffffffffL;
        }

        public String select(String rpcServiceKey) {
            byte[] digest = md5(rpcServiceKey);
            return selectForKey(hash(digest, 0));
        }

        public String selectForKey(long hashCode) {
            Map.Entry<Long, String> entry = virtualInvokers.tailMap(hashCode, true).firstEntry();

            if (null == entry) {
                entry = virtualInvokers.firstEntry();
            }

            return entry.getValue();
        }
    }
}
